from sklearn import datasets  # sklearn的数据集
from sklearn.tree import DecisionTreeClassifier,plot_tree  # sklearn模块的决策树分类器类
from matplotlib import pyplot as plt  # 绘图库

import numpy as np  # 矩阵运算库numpy
np.random.seed(0)

iris = datasets.load_iris()
# 变量
iris_x = iris.data
# 结果：有3个类别，分别是0， 1， 2
iris_y = iris.target

# 生成一个不重复的随机数组，元素值取[0~数据集长度]
random_arr = np.random.permutation(len(iris_x))

# 随机数组作为索引，划分训练集和测试集
iris_x_train = iris_x[random_arr[:-10]]
iris_y_train = iris_y[random_arr[:-10]]
iris_x_test = iris_x[random_arr[-10:]]
iris_y_test = iris_y[random_arr[-10:]]


clf = DecisionTreeClassifier(max_depth=4)
clf.fit(iris_x_train, iris_y_train)

iris_y_predict = clf.predict(iris_x_test)
score=clf.score(iris_x_test,iris_y_test,sample_weight=None)
print('预测结果 = ')
print(iris_y_predict)
print('实际值 = ')
print(iris_y_test)
print('得分:',score)


# ----------------------核心修改：定义中文特征名和类别名----------------------
# 中文特征名（对应原始特征：花萼长度、花萼宽度、花瓣长度、花瓣宽度）
cn_feature_names = ["花萼长度(cm)", "花萼宽度(cm)", "花瓣长度(cm)", "花瓣宽度(cm)"]

# 中文类别名（对应原始标签0、1、2，分别代表山鸢尾、变色鸢尾、维吉尼亚鸢尾）
cn_class_names = ["山鸢尾", "变色鸢尾", "维吉尼亚鸢尾"]

plt.figure(figsize=(12, 8))
# 设置中文字体，避免中文乱码
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
plot_tree(
    clf,
    feature_names=cn_feature_names,
    class_names=cn_class_names,
    filled=True,
    rounded=True
)
plt.show()
